"""
总调度器 - 负责协调各个专家代理的工作流程
"""

import asyncio
import time
from typing import Dict, List, Any, Optional, Tuple
from concurrent.futures import ThreadPoolExecutor, as_completed
from loguru import logger

from .config import Stage3Config
from .retriever import Retriever
from .utils.api_client import APIClient
from .utils.parallel_executor import ParallelExecutor
from .agents.logic_agent import LogicCoherenceAgent
from .agents.compliance_agent import ComplianceAgent
from .agents.cad_agent import CADReviewAgent
from .agents.table_agent import TableDataAgent
from .agents.format_agent import FormatAgent
from .agents.summary_agent import SummaryAgent


class MasterOrchestrator:
    """总调度器类，负责协调多个专家代理"""
    
    def __init__(self, config: Stage3Config):
        self.config = config
        self.retriever = Retriever(config)
        self.api_client = APIClient(config.api)
        self.parallel_executor = ParallelExecutor(
            max_workers=config.agent.max_parallel_agents,
            timeout=config.agent.agent_timeout
        )
        
        # 初始化专家代理
        self.agents = self._initialize_agents()
        
        logger.info(f"总调度器初始化完成，共{len(self.agents)}个专家代理")
    
    def _initialize_agents(self) -> Dict[str, Any]:
        """初始化所有专家代理"""
        agents = {}
        
        try:
            # 逻辑连贯性专家
            agents['logic_coherence'] = LogicCoherenceAgent(
                self.api_client, 
                self.config.agent.to_dict()
            )
            
            # 合规性专家
            agents['compliance'] = ComplianceAgent(
                self.api_client,
                self.config.agent.to_dict()
            )
            
            # CAD审查专家 - 临时禁用以避免处理时间过长
            # agents['cad_review'] = CADReviewAgent(
            #     self.api_client,
            #     self.config.agent.to_dict()
            # )
            logger.info("CAD审查代理已被临时禁用，跳过CAD图形分析")
            
            # 表格数据专家
            agents['table_data'] = TableDataAgent(
                self.api_client,
                self.config.agent.to_dict()
            )
            
            # 格式规范专家
            agents['format_review'] = FormatAgent(
                self.api_client,
                self.config.agent.to_dict()
            )
            
            # 总结专家
            agents['summary'] = SummaryAgent(
                self.api_client,
                self.config.agent.to_dict(),
                retriever=self.retriever
            )
            
            logger.info(f"成功初始化{len(agents)}个专家代理")
            return agents
            
        except Exception as e:
            logger.error(f"初始化专家代理失败: {e}")
            raise
    
    def process_session(self, session_id: str = None) -> Dict[str, Any]:
        """处理会话数据，协调所有专家代理进行审查"""
        start_time = time.time()
        
        try:
            # 获取会话摘要
            session_summary = self.retriever.get_session_summary(session_id)
            actual_session_id = session_summary['session_id']
            
            logger.info(f"开始处理会话: {actual_session_id}")
            logger.info(f"数据摘要: {session_summary.get('data_summary', {})}")
            
            # 检查是否启用并行处理
            if self.config.agent.enable_parallel:
                results = self._process_parallel(actual_session_id)
            else:
                results = self._process_sequential(actual_session_id)
            
            # 聚合结果
            aggregated_results = self._aggregate_results(results, session_summary)
            
            # 生成最终总结报告
            final_summary = self._generate_final_summary(results, aggregated_results, actual_session_id)
            
            processing_time = time.time() - start_time
            logger.info(f"会话处理完成，总耗时: {processing_time:.2f}秒")
            
            # 提取文本格式报告
            formatted_report = ""
            if final_summary.get('success', False) and 'formatted_report' in final_summary:
                formatted_report = final_summary['formatted_report']
            else:
                # 生成备用文本报告
                formatted_report = self._generate_simple_text_report(results, aggregated_results, processing_time)
            
            return {
                "success": True,
                "session_id": actual_session_id,
                "processing_time": processing_time,
                "session_summary": session_summary,
                "agent_results": results,
                "aggregated_results": aggregated_results,
                "final_summary": final_summary,
                "formatted_report": formatted_report,  # 添加文本格式报告
                "metadata": {
                    "processing_mode": "parallel" if self.config.agent.enable_parallel else "sequential",
                    "agents_count": len(self.agents),
                    "timestamp": time.time()
                }
            }
            
        except Exception as e:
            logger.error(f"会话处理失败: {e}")
            return {
                "success": False,
                "error": str(e),
                "session_id": session_id,
                "processing_time": time.time() - start_time,
                "metadata": {"timestamp": time.time()}
            }
    
    def _process_parallel(self, session_id: str) -> Dict[str, Any]:
        """并行处理所有专家代理（跳过总结代理）"""
        logger.info("使用并行模式处理专家代理任务")
        
        # 准备任务（跳过总结代理）
        tasks = []
        for agent_name, agent in self.agents.items():
            # 跳过总结代理，它将在所有专家代理完成后单独处理
            if agent_name == 'summary':
                continue
                
            task = {
                'task_id': f'agent_{agent_name}',
                'agent_name': agent_name,
                'agent': agent,
                'session_id': session_id
            }
            tasks.append(task)
        
        # 定义任务执行函数
        def execute_agent_task(task_data):
            """执行单个代理任务"""
            agent_name = task_data['agent_name']
            agent = task_data['agent']
            session_id = task_data['session_id']
            return self._process_single_agent(agent_name, agent, session_id)
        
        # 执行并行任务
        task_results = self.parallel_executor.execute_parallel(tasks, execute_agent_task)
        
        # 整理结果
        agent_results = {}
        for task_result in task_results:
            if task_result.success:
                agent_results[task_result.agent_name] = task_result.result
            else:
                agent_results[task_result.agent_name] = {
                    "success": False,
                    "error": task_result.error,
                    "agent_name": task_result.agent_name,
                    "processing_time": task_result.execution_time
                }
        
        return agent_results
    
    def _process_sequential(self, session_id: str) -> Dict[str, Any]:
        """顺序处理所有专家代理（跳过总结代理）"""
        logger.info("使用顺序模式处理专家代理任务")
        
        agent_results = {}
        
        for agent_name, agent in self.agents.items():
            # 跳过总结代理，它将在所有专家代理完成后单独处理
            if agent_name == 'summary':
                continue
                
            try:
                logger.info(f"开始处理代理: {agent_name}")
                result = self._process_single_agent(agent_name, agent, session_id)
                agent_results[agent_name] = result
                logger.info(f"代理 {agent_name} 处理完成")
                
            except Exception as e:
                logger.error(f"代理 {agent_name} 处理失败: {e}")
                agent_results[agent_name] = {
                    "success": False,
                    "error": str(e),
                    "agent_name": agent_name
                }
        
        return agent_results
    
    def _process_single_agent(self, agent_name: str, agent: Any, session_id: str) -> Dict[str, Any]:
        """处理单个代理"""
        start_time = time.time()
        
        try:
            # 为代理检索相关数据
            retrieval_result = self.retriever.retrieve_for_agent(agent.agent_type, session_id)
            
            if not retrieval_result.get('success', False):
                return {
                    "success": False,
                    "error": f"数据检索失败: {retrieval_result.get('error', '未知错误')}",
                    "agent_name": agent_name,
                    "agent_type": agent.agent_type,
                    "processing_time": time.time() - start_time
                }
            
            # 处理数据
            data = retrieval_result.get('data', [])
            context = retrieval_result.get('context', {})
            
            if not data:
                logger.warning(f"代理 {agent_name} 没有可处理的数据")
                return {
                    "success": True,
                    "agent_name": agent_name,
                    "agent_type": agent.agent_type,
                    "findings": [],
                    "summary": "无数据需要处理",
                    "confidence": 1.0,
                    "processing_time": time.time() - start_time,
                    "metadata": {"data_count": 0}
                }
            
            # 将location信息添加到context中
            if 'locations' in retrieval_result:
                context['locations'] = retrieval_result['locations']
            
            # 调用代理处理
            agent_result = agent.process(data, context)
            
            # 添加代理信息
            result_dict = agent_result.to_dict()
            result_dict.update({
                "agent_name": agent_name,
                "agent_type": agent.agent_type,
                "data_count": len(data),
                "retrieval_metadata": retrieval_result.get('metadata', {}),
                "locations": retrieval_result.get('locations', '未知位置')
            })
            
            return result_dict
            
        except Exception as e:
            logger.error(f"代理 {agent_name} 处理异常: {e}")
            return {
                "success": False,
                "error": str(e),
                "agent_name": agent_name,
                "agent_type": getattr(agent, 'agent_type', 'unknown'),
                "processing_time": time.time() - start_time
            }
    
    def _aggregate_results(self, agent_results: Dict[str, Any], session_summary: Dict[str, Any]) -> Dict[str, Any]:
        """聚合所有代理的结果"""
        try:
            # 统计信息
            total_findings = 0
            critical_count = 0
            major_count = 0
            minor_count = 0
            info_count = 0
            
            successful_agents = 0
            failed_agents = 0
            
            # 按类别组织发现的问题
            findings_by_category = {}
            findings_by_severity = {
                'critical': [],
                'major': [],
                'minor': [],
                'info': []
            }
            
            # 代理摘要
            agent_summaries = {}
            
            for agent_name, result in agent_results.items():
                if result.get('success', False):
                    successful_agents += 1
                    
                    # 统计发现的问题
                    findings = result.get('findings', [])
                    total_findings += len(findings)
                    
                    for finding in findings:
                        severity = finding.get('severity', 'info')
                        category = finding.get('category', '其他')
                        
                        # 按严重性统计
                        if severity == 'critical':
                            critical_count += 1
                        elif severity == 'major':
                            major_count += 1
                        elif severity == 'minor':
                            minor_count += 1
                        else:
                            info_count += 1
                        
                        # 按严重性分组
                        findings_by_severity[severity].append({
                            **finding,
                            'agent_name': agent_name,
                            'agent_type': result.get('agent_type', '')
                        })
                        
                        # 按类别分组
                        if category not in findings_by_category:
                            findings_by_category[category] = []
                        findings_by_category[category].append({
                            **finding,
                            'agent_name': agent_name,
                            'agent_type': result.get('agent_type', '')
                        })
                    
                    # 代理摘要
                    agent_summaries[agent_name] = {
                        'agent_type': result.get('agent_type', ''),
                        'summary': result.get('summary', ''),
                        'confidence': result.get('confidence', 0.0),
                        'findings_count': len(findings),
                        'processing_time': result.get('processing_time', 0.0),
                        'data_count': result.get('data_count', 0)
                    }
                else:
                    failed_agents += 1
                    agent_summaries[agent_name] = {
                        'agent_type': result.get('agent_type', ''),
                        'error': result.get('error', '未知错误'),
                        'processing_time': result.get('processing_time', 0.0)
                    }
            
            # 计算整体评分
            overall_score = self._calculate_overall_score(
                critical_count, major_count, minor_count, total_findings
            )
            
            # 生成整体摘要
            overall_summary = self._generate_overall_summary(
                agent_summaries, total_findings, critical_count, major_count, minor_count
            )
            
            return {
                "statistics": {
                    "total_findings": total_findings,
                    "severity_distribution": {
                        "critical": critical_count,
                        "major": major_count,
                        "minor": minor_count,
                        "info": info_count
                    },
                    "agent_performance": {
                        "successful": successful_agents,
                        "failed": failed_agents,
                        "total": len(agent_results)
                    }
                },
                "findings_by_severity": findings_by_severity,
                "findings_by_category": findings_by_category,
                "agent_summaries": agent_summaries,
                "overall_assessment": {
                    "score": overall_score,
                    "summary": overall_summary,
                    "recommendation": self._generate_recommendation(overall_score, critical_count, major_count)
                },
                "session_info": session_summary
            }
            
        except Exception as e:
            logger.error(f"结果聚合失败: {e}")
            return {
                "error": f"结果聚合失败: {str(e)}",
                "raw_results": agent_results,
                "session_info": session_summary
            }
    
    def _calculate_overall_score(self, critical: int, major: int, minor: int, total: int) -> float:
        """计算整体评分"""
        if total == 0:
            return 10.0
        
        # 权重扣分
        deduction = critical * 3.0 + major * 2.0 + minor * 1.0
        
        # 基础分数
        base_score = 10.0
        
        # 根据问题数量调整扣分
        if total > 0:
            severity_ratio = (critical * 3 + major * 2 + minor) / total
            adjusted_deduction = deduction * (1 + severity_ratio * 0.1)
        else:
            adjusted_deduction = 0
        
        final_score = max(0.0, base_score - adjusted_deduction)
        return round(final_score, 2)
    
    def _generate_overall_summary(self, agent_summaries: Dict[str, Any], 
                                total_findings: int, critical: int, major: int, minor: int) -> str:
        """生成整体摘要"""
        summary_parts = []
        
        # 基本统计
        summary_parts.append(f"本次审查共发现{total_findings}个问题")
        
        if critical > 0:
            summary_parts.append(f"其中严重问题{critical}个")
        if major > 0:
            summary_parts.append(f"重要问题{major}个")
        if minor > 0:
            summary_parts.append(f"轻微问题{minor}个")
        
        # 代理表现
        successful_agents = [name for name, summary in agent_summaries.items() 
                           if 'error' not in summary]
        
        if successful_agents:
            summary_parts.append(f"成功完成审查的专家代理: {', '.join(successful_agents)}")
        
        # 主要问题领域
        if critical > 0:
            summary_parts.append("建议优先处理严重问题")
        elif major > 0:
            summary_parts.append("建议及时处理重要问题")
        elif minor > 0:
            summary_parts.append("建议逐步改进轻微问题")
        else:
            summary_parts.append("整体质量良好，无重大问题")
        
        return "；".join(summary_parts) + "。"
    
    def _generate_recommendation(self, score: float, critical: int, major: int) -> str:
        """生成改进建议"""
        if score >= 9.0:
            return "方案质量优秀，可以进入下一阶段"
        elif score >= 7.0:
            if critical > 0:
                return "存在严重问题，建议修正后重新审查"
            elif major > 0:
                return "存在重要问题，建议修正后继续"
            else:
                return "整体质量良好，建议优化细节"
        elif score >= 5.0:
            return "存在较多问题，建议全面修正后重新提交"
        else:
            return "问题较为严重，建议重新设计方案"
    
    def get_agent_status(self) -> Dict[str, Any]:
        """获取所有代理的状态"""
        status = {}
        
        for agent_name, agent in self.agents.items():
            status[agent_name] = {
                "name": agent.name,
                "agent_type": agent.agent_type,
                "initialized": True,
                "config": {
                    "timeout": self.config.agent.agent_timeout,
                    "weight": self.config.agent.get_agent_weight(agent_name)
                }
            }
        
        return {
            "total_agents": len(self.agents),
            "parallel_enabled": self.config.agent.enable_parallel,
            "max_parallel": self.config.agent.max_parallel_agents,
            "agents": status
        }
    
    def _generate_final_summary(self, agent_results: Dict[str, Any], aggregated_results: Dict[str, Any], session_id: str) -> Dict[str, Any]:
        """生成最终总结报告"""
        start_time = time.time()
        
        try:
            logger.info("开始生成最终总结报告")
            
            # 获取总结代理
            summary_agent = self.agents.get('summary')
            if not summary_agent:
                logger.warning("总结代理未初始化，跳过最终总结")
                return {
                    "success": False,
                    "error": "总结代理未初始化",
                    "summary": "无法生成最终总结报告",
                    "processing_time": time.time() - start_time,
                    "formatted_report": "错误：总结代理未初始化，无法生成报告。"
                }
            
            # 准备专家代理结果列表
            expert_results = []
            for agent_name, result in agent_results.items():
                # 跳过总结代理本身（虽然现在它不在agent_results中）
                if agent_name == 'summary':
                    continue
                    
                # 转换为AgentResult对象格式
                if result.get('success', False):
                    from .agents.base_agent import AgentResult
                    
                    # 直接创建AgentResult对象
                    agent_result = AgentResult(
                        agent_name=result.get('agent_name', agent_name),
                        agent_type=result.get('agent_type', 'unknown'),
                        success=result.get('success', False),
                        findings=result.get('findings', []),
                        summary=result.get('summary', ''),
                        confidence=result.get('confidence', 0.0),
                        processing_time=result.get('processing_time', 0.0),
                        metadata=result.get('metadata', {}),
                        error=result.get('error')
                    )
                    expert_results.append(agent_result)
            
            if not expert_results:
                logger.warning("没有成功的专家代理结果用于总结")
                return {
                    "success": False,
                    "error": "没有可用的专家代理结果",
                    "summary": "无专家代理结果可供总结",
                    "processing_time": time.time() - start_time,
                    "formatted_report": "错误：没有可用的专家代理结果，无法生成报告。"
                }
            
            # 准备上下文信息
            context = {
                "session_id": session_id,
                "aggregated_statistics": aggregated_results.get('statistics', {}),
                "session_info": aggregated_results.get('session_info', {}),
                "overall_assessment": aggregated_results.get('overall_assessment', {})
            }
            
            logger.info(f"调用总结代理处理 {len(expert_results)} 个专家代理结果")
            
            # 直接调用总结代理的process方法，传入专家代理结果列表
            summary_result = summary_agent.process(expert_results, context)
            
            logger.info("最终总结报告生成完成")
            
            processing_time = time.time() - start_time
            
            # 提取Markdown格式报告
            formatted_report = ""
            if summary_result.success and hasattr(summary_result, 'metadata') and summary_result.metadata:
                # 优先使用markdown_report，如果没有则使用formatted_report作为兼容
                formatted_report = summary_result.metadata.get('markdown_report', '') or summary_result.metadata.get('formatted_report', '')
            
            # 如果没有Markdown报告，生成一个基本的报告
            if not formatted_report:
                formatted_report = self._generate_fallback_text_report(expert_results, summary_result)
            
            return {
                "success": True,
                "summary_result": summary_result.to_dict() if hasattr(summary_result, 'to_dict') else summary_result,
                "expert_results_count": len(expert_results),
                "processing_time": processing_time,
                "formatted_report": formatted_report  # 添加文本格式报告
            }
            
        except Exception as e:
            logger.error(f"生成最终总结报告失败: {e}")
            fallback_report = f"""错误报告
========================
生成时间: {time.strftime('%Y-%m-%d %H:%M:%S')}
错误信息: {str(e)}
处理时间: {time.time() - start_time:.2f}秒

请检查系统日志获取更多详细信息。
========================"""
            
            return {
                "success": False,
                "error": str(e),
                "summary": "最终总结报告生成失败",
                "processing_time": time.time() - start_time,
                "formatted_report": fallback_report
            }
    
    def _generate_fallback_text_report(self, expert_results: List[Any], summary_result: Any) -> str:
        """生成备用文本报告"""
        try:
            report_lines = []
            report_lines.append("设计方案审查报告")
            report_lines.append("=" * 50)
            report_lines.append(f"生成时间: {time.strftime('%Y-%m-%d %H:%M:%S')}")
            report_lines.append(f"专家代理数量: {len(expert_results)}")
            report_lines.append("")
            
            # 添加总结信息
            if hasattr(summary_result, 'summary') and summary_result.summary:
                report_lines.append("总体评估:")
                report_lines.append("-" * 20)
                report_lines.append(summary_result.summary)
                report_lines.append("")
            
            # 添加专家代理结果
            report_lines.append("专家代理审查结果:")
            report_lines.append("-" * 30)
            
            for result in expert_results:
                agent_name = getattr(result, 'agent_name', '未知代理')
                agent_summary = getattr(result, 'summary', '无摘要')
                findings_count = len(getattr(result, 'findings', []))
                
                report_lines.append(f"• {agent_name}: {agent_summary} (发现 {findings_count} 个问题)")
            
            report_lines.append("")
            report_lines.append("=" * 50)
            
            return "\n".join(report_lines)
            
        except Exception as e:
            return f"生成备用报告时出错: {str(e)}"
    
    def _generate_simple_text_report(self, agent_results: Dict[str, Any], aggregated_results: Dict[str, Any], processing_time: float) -> str:
        """生成简单的文本报告"""
        try:
            report_lines = []
            report_lines.append("设计方案审查报告")
            report_lines.append("=" * 50)
            report_lines.append(f"生成时间: {time.strftime('%Y-%m-%d %H:%M:%S')}")
            report_lines.append(f"处理时间: {processing_time:.2f}秒")
            report_lines.append("")
            
            # 统计信息
            stats = aggregated_results.get('statistics', {})
            if stats:
                report_lines.append("统计信息:")
                report_lines.append("-" * 20)
                report_lines.append(f"总问题数: {stats.get('total_findings', 0)}")
                
                severity_dist = stats.get('severity_distribution', {})
                if severity_dist:
                    report_lines.append(f"严重问题: {severity_dist.get('critical', 0)}")
                    report_lines.append(f"重要问题: {severity_dist.get('major', 0)}")
                    report_lines.append(f"轻微问题: {severity_dist.get('minor', 0)}")
                    report_lines.append(f"信息问题: {severity_dist.get('info', 0)}")
                
                agent_perf = stats.get('agent_performance', {})
                if agent_perf:
                    report_lines.append(f"成功代理: {agent_perf.get('successful', 0)}")
                    report_lines.append(f"失败代理: {agent_perf.get('failed', 0)}")
                
                report_lines.append("")
            
            # 整体评估
            overall = aggregated_results.get('overall_assessment', {})
            if overall:
                report_lines.append("整体评估:")
                report_lines.append("-" * 20)
                report_lines.append(f"评分: {overall.get('score', 0)}/10")
                report_lines.append(f"摘要: {overall.get('summary', '无摘要')}")
                report_lines.append(f"建议: {overall.get('recommendation', '无建议')}")
                report_lines.append("")
            
            # 代理结果摘要
            report_lines.append("代理审查结果:")
            report_lines.append("-" * 30)
            
            for agent_name, result in agent_results.items():
                if result.get('success', False):
                    findings_count = len(result.get('findings', []))
                    summary = result.get('summary', '无摘要')
                    report_lines.append(f"• {agent_name}: {summary} (发现 {findings_count} 个问题)")
                else:
                    error = result.get('error', '未知错误')
                    report_lines.append(f"• {agent_name}: 处理失败 - {error}")
            
            report_lines.append("")
            report_lines.append("=" * 50)
            
            return "\n".join(report_lines)
            
        except Exception as e:
            return f"生成简单文本报告时出错: {str(e)}"
    
    async def process_session_async(self, session_id: str = None) -> Dict[str, Any]:
        """异步处理会话（如果需要的话）"""
        # 将同步方法包装为异步
        loop = asyncio.get_event_loop()
        return await loop.run_in_executor(None, self.process_session, session_id)